
"""
参考：miner-vl:  tests/test_layout.py
"""
from typing import List
from PIL import Image

from .sglang_client_predictor import SglangClientPredictor
from .vlm_magic_model import MagicModel
from .vlm_middle_json_mkcontent import mk_blocks_to_markdown
from .enum_class import MakeMode
from ..config import settings
from ..utils.utils import image_to_b64str


class VlmAnalyzer(object):

    def __init__(self, server_url: str):
        self.predictor = SglangClientPredictor(server_url)

    def batch_analyze(self, images: List[Image.Image], prompts: List[str] | str = ''):
        """表格识别"""
        batch_base64_images = [image_to_b64str(image) for image in images]
        token_list = self.predictor.batch_predict(
            images=batch_base64_images,
            prompts=prompts
        )
        page_markdowns = []
        for image, token in zip(images, token_list):
            width, height = image.size
            page_blocks = self.token_to_image_info(token, width, height)
            block_markdowns = mk_blocks_to_markdown(
                page_blocks, MakeMode.MM_MD, formula_enable=True, table_enable=True, img_buket_path=''
            )
            page_markdown = '\n\n'.join(block_markdowns)
            page_markdowns.append(page_markdown)
        return page_markdowns

    @staticmethod
    def token_to_image_info(token, width, height):
        magic_model = MagicModel(token, width, height)
        image_blocks = magic_model.get_image_blocks()
        table_blocks = magic_model.get_table_blocks()
        title_blocks = magic_model.get_title_blocks()
        text_blocks = magic_model.get_text_blocks()
        interline_equation_blocks = magic_model.get_interline_equation_blocks()
        page_blocks = []
        page_blocks.extend([*image_blocks, *table_blocks, *title_blocks, *text_blocks, *interline_equation_blocks])
        page_blocks.sort(key=lambda x: x["index"])
        return page_blocks

    def batch_flows_analyze(self, batch_base64_images: List):
        """流程图识别，只能走Qwen-Vl了 """
        results = self.predictor.batch_predict(
            images=batch_base64_images,
            prompts='Flow Chart Parsing:'
        )


if __name__ == '__main__':
    vlm = VlmAnalyzer(settings.VLM_SERVER_URL)
    pil_img = Image.open(r'D:\workprojects\marker\data\images\table_002.png')
    # pil_img = Image.open(r'D:\workprojects\marker\0.jpg')
    vlm.batch_tables_analyze([pil_img])